#include "parser.h"
#include "translator.h"
#include "cfg.h"

#define _CFG 1
#define _TAC 2
#define _C	3
#define _REP 4

#define _NOPT	0
#define _SCP 1
#define _DSE 2

void IR2C(IR& ir, ostream& o){
	Translator translator(ir.instructions);
	translator.translate();
	translator.output(o);
}

int spills = 0;
vector<string> spilled_vars;

void toAssembly(IR& ir, vector<CFG *>& cfg, ostream& o, bool report){
	vector<Instr>& is = ir.instructions;
	for(int i = 0; i < is.size(); i++){
		if(is[i].opt == CALL)
			continue;

		if(is[i].param1.type == LABEL){
			is[is[i].param1.value].hasLabel = true;
		}

		if(is[i].param2.type == LABEL){
			is[is[i].param2.value].hasLabel = true;
		}
	}

	vector<string> assembly;

	assembly.push_back(".global gp\n");
	assembly.push_back(".data\n");
	assembly.push_back(".comm gp, 32768, 32\n");
	assembly.push_back(".LC0:\n");
	assembly.push_back(".string \" %d\"\n");
	assembly.push_back(".text\n");
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->toAssembly(o, assembly);
	}

	if(!report){
		for(int i = 0; i < assembly.size(); i++){
			o << assembly[i];
		}
	}
	else{
		cout << "Number of spills: " << spills << endl;
		cout << "Number of assembly instructions generated: " << assembly.size() << endl;
	}
}

void buildCFG(IR&ir, vector<CFG *>& cfg){
	int start = 1;
	int end;
	while(start < ir.instructions.size()){
		CFG * pCFG = new CFG(ir.instructions);
		if(pCFG->build(start) == false){
			delete pCFG;
			return;
		}
		cfg.push_back(pCFG);
		start = pCFG->end + 1;
	}
}

void deleteCFG(vector<CFG *>& cfg){
	for(int i = 0 ; i < cfg.size(); i++){
		delete cfg[i];
	}
	cfg.clear();
}

void dse(IR&ir,vector<CFG *>& cfg){
	for(int i = 0; i < cfg.size(); i++){
		do{
			cfg[i]->liveness();
		}while(cfg[i]->dse());
	}
	ir.refine();
	deleteCFG(cfg);
	buildCFG(ir, cfg);
}

void printCFG(vector<CFG *>& cfg, ostream& o){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->output(o);
	}
}

void buildSSA(IR&ir, vector<CFG *>& cfg){
	vector<Instr> oldir = ir.instructions;
	map<int, int> newreg;
	map<int, int> newlab;
	map<int, BasicBlock *> bb;

	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->findDominators();
		cfg[i]->buildDomTree();
		cfg[i]->findDomFrontier(cfg[i]->nodes.begin()->second);
		cfg[i]->liveness();
		//cfg[i]->printDomFrontier(cout);
		cfg[i]->placePhi();
		//cfg[i]->printPhi(cout);
		bb.insert(cfg[i]->nodes.begin(), cfg[i]->nodes.end());
	}

	ir.instructions.clear();
	int c = 0;
	for(int i = 0; i < oldir.size(); i++){
		if(bb.find(i) != bb.end()){
			BasicBlock * b = bb[i];
			newlab[i] = c;
			for(int j = 0 ; j < b->phis.size(); j++){//insert phi funcs
				ir.instructions.push_back(b->phis[j]);
			}
			c += b->phis.size();
		}
		ir.instructions.push_back(oldir[i]);
		newreg[i] = c;
		c++;
	}

	//update r* label*
	for(int i = 0; i < ir.instructions.size(); i++){
		Instr& instr = ir.instructions[i];
		if(instr.param1.type == REG)
			instr.param1.value = newreg[instr.param1.value];
		else if(instr.param1.type == LABEL)
			instr.param1.value = newlab[instr.param1.value];

		if(instr.param2.type == REG)
			instr.param2.value = newreg[instr.param2.value];
		else if(instr.param2.type == LABEL)
			instr.param2.value = newlab[instr.param2.value];

		if(instr.opt == MOVE && instr.param1.type == REG && instr.param1.value == 0){
			instr.param1.value = i - 1;
		}
	}

	//rebuild cfg
	deleteCFG(cfg);
	buildCFG(ir, cfg);

	//now do renaming
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->findDominators();
		cfg[i]->buildDomTree();
		cfg[i]->rename();
	}
	
}

void removePhi(IR&ir, vector<CFG*>& cfg){
	vector<Instr> oldir = ir.instructions;
	map<int, int> newreg;
	map<int, int> newlab;
	map<int, BasicBlock *> bb;

	for(int i = 0; i < cfg.size(); i++){
		bb.insert(cfg[i]->nodes.begin(), cfg[i]->nodes.end());
	}

	ir.instructions.clear();
	int c = 0;
	for(int i = 0; i < oldir.size(); i++){
		if(bb.find(i) != bb.end()){
			BasicBlock * b = bb[i];
			newlab[i] = c;
			int nphi = 0;
			for(int j = b->in; j < b->out && oldir[j].opt==PHI; j+=2){
				nphi++;
			}
			i += (2*nphi);
		}
		ir.instructions.push_back(oldir[i]);
		newreg[i] = c;
		c++;
	}

	//update r* label*
	for(int i = 0; i < ir.instructions.size(); i++){
		Instr& instr = ir.instructions[i];
		if(instr.param1.type == REG)
			instr.param1.value = newreg[instr.param1.value];
		else if(instr.param1.type == LABEL)
			instr.param1.value = newlab[instr.param1.value];
		else if(instr.param1.type == LOCAL_VAR)
			instr.param1.var_id = -1;

		if(instr.param2.type == REG)
			instr.param2.value = newreg[instr.param2.value];
		else if(instr.param2.type == LABEL)
			instr.param2.value = newlab[instr.param2.value];
		else if(instr.param2.type == LOCAL_VAR)
			instr.param2.var_id = -1;

	}
	//rebuild cfg
	deleteCFG(cfg);
	buildCFG(ir, cfg);
}


void dseReport(vector<CFG *>& cfg, ostream& o){
	for(int i = 0; i < cfg.size(); i++){
		o<<"Function: "<< cfg[i]->entry<<endl;
		o<<"Number of statements eliminated in SCR: "<< cfg[i]->dseInSCR<<endl;
		o<<"Number of statements eliminated not in SCR: "<<cfg[i]->dseOutSCR<<endl;
	}
}

void findSCR(vector<CFG *>& cfg){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->findSCR();
	}
}

void printSCR(vector<CFG *>& cfg, ostream& o){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->printSCR(o);
	}
}

void scp(vector<CFG *>& cfg){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->scp();
	}
}

void scpReport(vector<CFG *>& cfg, ostream& o){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->printSCP(o);
	}
}

//SSA Constant Propagation
void ssacp(vector<CFG *>& cfg){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->ssacp();
	}
}

//SSA Loop Invarient Code Motion
void ssalicm(IR& ir, vector<CFG *>& cfg){
	for(int i = 0; i < cfg.size(); i++){
		vector<vector<int> >* SCR = &cfg[i]->SCR;

		if(!cfg[i]->scr_found){
			cfg[i]->findSCR();
		}
		if(SCR->size() == 0)
			continue;

		bool changed = false;
		int scrsize = SCR->size();

		for(int j = 0; j < scrsize;){
			changed = cfg[i]->ssalicm(j);
			if(changed){
				int lines_moved = cfg[i]->lines_moved;
				int num_propagated = cfg[i]->num_propagated;
				//cout << "lines_moved: " << lines_moved << endl;
				//ir.print(cout);
				deleteCFG(cfg);
				buildCFG(ir, cfg);
				cfg[i]->findSCR();
				SCR = &cfg[i]->SCR;
				//printSCR(cfg, cout);
				//printCFG(cfg, cout);
				cfg[i]->lines_moved = lines_moved;
				cfg[i]->num_propagated = num_propagated;
			}
			else
				j++;
		}

		int lines_moved = cfg[i]->lines_moved;
		int num_propagated = cfg[i]->num_propagated;
		//printSCR(cfg, cout);
		//printCFG(cfg, cout);
		deleteCFG(cfg);
		buildCFG(ir, cfg);
		cfg[i]->findSCR();
		cfg[i]->lines_moved = lines_moved;
		cfg[i]->num_propagated = num_propagated;
	}
}

void ssalicmReport(vector<CFG *>& cfg, ostream& o){
	for(int i = 0; i < cfg.size(); i++){
		cfg[i]->printssalicm(o);
	}
}

void registerAllocation(IR& ir, vector<CFG *>&  cfg);

int main(int argc, char * argv[]){
	
	//ostream& o = cout;
	//ifstream ifs("C:\\Users\\rliu\\Documents\\Visual Studio 2008\\Projects\\cs380c\\examples\\prime");
	//IR ir(cin);
	//vector<CFG *> vCFG;
	//buildCFG(ir, vCFG);
	//buildSSA(ir, vCFG);
	//ir.print(o);
	//findSCR(vCFG);
	//printSCR(vCFG, o);
	//ssacp(vCFG);
	//scpReport(vCFG, o);
	//ssalicm(ir, vCFG);
	//ssalicmReport(vCFG, o);
	//ir.print(o);
	//o<<"************************"<<endl;
	//removePhi(ir, vCFG);
	//ir.print(o);
	//deleteCFG(vCFG);

	ostream& o = cout;

	IR ir(cin);
	vector<CFG *> vCFG;
	buildCFG(ir, vCFG);
	findSCR(vCFG);
	dse(ir,vCFG);
	//ir.print(cout);
	bool report = false;
	for(int i = 1; i < argc; i++){
		string s = argv[i]; 
		if(s.find("-a") != string::npos){
			report = false;
		}
		if(s.find("-r") != string::npos){
			report = true;
		}
	}

	registerAllocation(ir, vCFG);
	//ir.print(o);
	toAssembly(ir, vCFG, cout, report);
}

void refactorLineNumbers(IR& ir, int line_added, bool jump_to){
	vector<Instr>& is = ir.instructions;
	for(int i = 0; i < is.size(); i++){
		if( (is[i].param1.type == REG) && (is[i].param1.value >= line_added) && (i != line_added + 1)){
			is[i].param1.value++;
		}

		if( (is[i].param2.type == REG) && (is[i].param2.value >= line_added) && (i != line_added + 1)){
			is[i].param2.value++;
		}

		if(jump_to){
			if( (is[i].param1.type == LABEL) && (is[i].param1.value > line_added) ){
				is[i].param1.value++;
			}
	
			if( (is[i].param2.type == LABEL) && (is[i].param2.value > line_added) ){
				is[i].param2.value++;
			}
		}
		else{
			if( (is[i].param1.type == LABEL) && (is[i].param1.value >= line_added) ){
				is[i].param1.value++;
			}
	
			if( (is[i].param2.type == LABEL) && (is[i].param2.value >= line_added) ){
				is[i].param2.value++;
			}
		}
	}
}

void registerAllocation(IR& ir, vector<CFG *>& vCFG){
	int spill_op;
	int spill_op_interference;
	int size;
	bool reg;
	string spill_var;
	map<int, BasicBlock *>::iterator it;
	BasicBlock * b;
	set<string>::iterator dit;
	vector<Instr>& is = ir.instructions;

	//ir.print(cout);

	int j, k;

	for(int x = 0; x < vCFG.size(); x++){
		CFG* cfg = vCFG[x];
		cfg->makeInterferenceGraph();

		while(!cfg->colorInterferenceGraph()){
			spill_var = cfg->spilled[0];
			spills++;
			spilled_vars.push_back(spill_var);
			bool is_reg;
			int offset;
			if(atoi(spill_var.c_str()) != 0){
				is_reg = true;
				is[cfg->entry].param1.value += 8;
				offset = 0 - is[cfg->entry].param1.value;
			}
			else
				is_reg = false;

			int begin = cfg->entry;
			int end = cfg->end;

			for(int i = begin; i <= end; i++){
				Instr& cur_inst = is[i];
				//if(cur_inst.spill == true)
				//	continue;

				
				Instr new_inst;
				new_inst.spill = true;
				bool is_move_inst;
				bool do_insert = false;
				if(cur_inst.opt != MOVE && (!is_reg || atoi(spill_var.c_str()) != i)){
					is_move_inst = false;
					new_inst.sOpt = "load";
					new_inst.opt = LOAD;
					new_inst.param1.type = FP;
					new_inst.param1.literal = "";
				}
				else{
					is_move_inst = true; 
				}


				if(is_reg == false){ 
					if(cur_inst.param1.type == LOCAL_VAR && cur_inst.param1.literal.compare(spill_var) == 0){
						if(!is_move_inst){
							new_inst.param1.value = cur_inst.param1.value;
							cur_inst.param1.type = REG;
							cur_inst.param1.value = i;
							cur_inst.param1.literal = "";
						}
						else{
						}
						do_insert = true;
					}

					if(cur_inst.param2.type == LOCAL_VAR && cur_inst.param2.literal.compare(spill_var) == 0){
						if(!is_move_inst){
							new_inst.param1.value = cur_inst.param2.value;
							cur_inst.param2.type = REG;
							cur_inst.param2.value = i;
							cur_inst.param2.literal = "";
						}
						else{
						}
						do_insert = true;
					}
				}
				else if(is_reg == true){
					if(atoi(spill_var.c_str()) == i){
						new_inst.param2.value = offset;
						do_insert = true;
					}
 					else{
						if((cur_inst.param1.type == REG) && (cur_inst.param1.value == atoi(spill_var.c_str()))){
							new_inst.param1.value = offset;
							cur_inst.param1.value = i;
							do_insert = true;
						}
	
						if((cur_inst.param2.type == REG) && (cur_inst.param2.value == atoi(spill_var.c_str()))){
							new_inst.param1.value = offset;
							cur_inst.param2.value = i;
							do_insert = true;
						}
					}
				}

				if(do_insert){
					if(!is_move_inst){
						is.insert(is.begin() + i, new_inst);
						refactorLineNumbers(ir, i, true);
						end++;
					}
					else{
						cur_inst.sOpt = "store";
						cur_inst.opt = STORE;
						cur_inst.param2.type = FP;
						cur_inst.param2.literal = "";
						if(is_reg)
							cur_inst.param2.value = offset;
					}
				}
				//cout << endl << "spill_var: " << spill_var << " i: " << i << endl;
				//ir.print(cout);
			}

			vector<map<string, int> > saved_var2reg;
			for(int m = 0; m < vCFG.size(); m++){
				saved_var2reg.push_back(vCFG[m]->var2reg);
			}
			deleteCFG(vCFG);
			buildCFG(ir, vCFG);
			for(int m = 0; m < vCFG.size(); m++){
				vCFG[m]->var2reg = saved_var2reg[m];
			}
			cfg = vCFG[x];
			cfg->makeInterferenceGraph();
		}

		cfg->var2reg = cfg->coloring;

	}
}

